{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import hydra\n",
    "import torch\n",
    "import functools\n",
    "import numpy as np \n",
    "import yaml\n",
    "import sys, os\n",
    "sys.path.append(\"/home/user/sirf/\")\n",
    "from src import (BrainWebOSEM, kl_div)\n",
    "from omegaconf import DictConfig, OmegaConf\n",
    "import time\n",
    "import matplotlib.pyplot as plt\n",
    "from src import PSNR, SSIM\n",
    "\n",
    "import cupy as xp\n",
    "\n",
    "detector_efficiency = 1./30\n",
    "\n",
    "def get_acq_model():\n",
    "\timport pyparallelproj.coincidences as coincidences\n",
    "\timport pyparallelproj.petprojectors as petprojectors\n",
    "\timport pyparallelproj.resolution_models as resolution_models\n",
    "\timport cupyx.scipy.ndimage as ndi\n",
    "\t\"\"\"\n",
    "\tcreate forward operator\n",
    "\t\"\"\"\n",
    "\tcoincidence_descriptor = coincidences.GEDiscoveryMICoincidenceDescriptor(\n",
    "\t\tnum_rings=1,\n",
    "\t\tsinogram_spatial_axis_order=coincidences.SinogramSpatialAxisOrder['RVP'],xp=xp)\n",
    "\tacq_model = petprojectors.PETJosephProjector(coincidence_descriptor,\n",
    "\t\t(128, 128, 1), (-127.0, -127.0, 0.0), (2., 2., 2.))\n",
    "\tres_model = resolution_models.GaussianImageBasedResolutionModel(\n",
    "\t\t(128, 128, 1), tuple(4.5 / (2.35 * x) for x in (2., 2., 2.)), xp, ndi)\n",
    "\tacq_model.image_based_resolution_model = res_model\n",
    "\treturn acq_model\n",
    "\n",
    "\n",
    "# Generate a unique filename, the file directories specify the SDE and data\n",
    "name = \"\"\n",
    "results = {}\n",
    "\n",
    "timestr = time.strftime(\"%Y%m%d_%H%M%S_\")\n",
    "\n",
    "###### SET SEED ######\n",
    "torch.manual_seed(42)\n",
    "np.random.seed(42)\n",
    "\n",
    "###### GET ACQUISITION MODEL AND DATA ######\n",
    "# get the acquisition model\n",
    "acq_model = get_acq_model()\n",
    "dataset = BrainWebOSEM(part=\"test\",\n",
    "        noise_level=\"2.5\", \n",
    "        base_path=\"/home/user/sirf/src/brainweb_2d/\")\n",
    "subset = list(range(2, len(dataset), 4))\n",
    "dataset = torch.utils.data.Subset(dataset, subset)\n",
    "test_loader = torch.utils.data.DataLoader(dataset, \n",
    "    batch_size=1, shuffle=False)\n",
    "# as there are 10 realisations then batch = 10\n",
    "batch_size = 10\n",
    "print(\"Length of test loader: \", len(test_loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "from glob import glob\n",
    "import torch\n",
    "import omegaconf\n",
    "from normalisation import Normalisation\n",
    "from lpd_modules import LPDForwardFunction2D\n",
    "from lpd import get_lpd_model\n",
    "from unet import get_unet_model\n",
    "from src import PSNR, SSIM\n",
    "\n",
    "detector_efficiency = 1./30\n",
    "\n",
    "models = [\"data_corrected_mean\"]\n",
    "\n",
    "model_base_path = \"/home/user/sirf/results/lpd_unet/UNET/\"\n",
    "save_idxs = [10,20,30,40,50,60]\n",
    "for model_type in models:\n",
    "    model_path = model_base_path + model_type + \"/\"\n",
    "    config = omegaconf.OmegaConf.load(model_path + \".hydra/config.yaml\")\n",
    "    model = get_unet_model(in_ch=config.benchmark.in_ch, \n",
    "                            out_ch=config.benchmark.out_ch, \n",
    "                            scales=config.benchmark.scales, \n",
    "                            skip=config.benchmark.skip,\n",
    "                            channels=config.benchmark.channels, \n",
    "                            use_sigmoid=config.benchmark.use_sigmoid,\n",
    "                            use_norm=config.benchmark.use_norm)\n",
    "    model.load_state_dict(torch.load(sorted(glob(model_path+\"/*/model_min_val_loss.pt\"))[-1]))\n",
    "    model.eval()\n",
    "    model.to(config.device)\n",
    "    print(sum(p.numel() for p in model.parameters()))\n",
    "    get_normalisation = Normalisation(config.benchmark.normalisation)\n",
    "    stdf = []\n",
    "    crcf = []\n",
    "    psnrf = []\n",
    "    ssimf = []\n",
    "    kl_divf = []\n",
    "    for idx, batch in enumerate(test_loader):\n",
    "        ref = torch.swapaxes(batch[0], 0, 1).to(config.device)\n",
    "        osem = torch.swapaxes(batch[2], 0, 1).to(config.device)\n",
    "        measurements = torch.swapaxes(batch[4], 0, 1).to(config.device)\n",
    "        contamination_factor = torch.swapaxes(batch[5], 0, 1)[:,[0],None].to(config.device)[:,0,:]\n",
    "        attn_factors = torch.swapaxes(batch[6], 0, 1).to(config.device)\n",
    "        if len(batch) > 7:\n",
    "            background = torch.swapaxes(batch[7], 0, 1).to(config.device)\n",
    "            tumour_rois = torch.swapaxes(batch[8], 0, 1).to(config.device)\n",
    "        norm = get_normalisation(osem, measurements, contamination_factor)\n",
    "        x_pred = torch.clamp(model(osem, norm).detach(),0)\n",
    "        kldiv_r = kl_div(x = x_pred,\n",
    "\t\t\tacq_model=acq_model, \n",
    "\t\t\tattn_factors=attn_factors, \n",
    "\t\t\tcontamination=contamination_factor.unsqueeze(1), \n",
    "\t\t\tmeasurements=measurements,\n",
    "\t\t\tscale_factor=1.)[0].squeeze()\n",
    "        kl_divf.append(kldiv_r.mean())\n",
    "        \n",
    "        images = x_pred.squeeze().unsqueeze(0)\n",
    "        refs = ref.squeeze().unsqueeze(0)\n",
    "        if len(batch) > 7:\n",
    "            lesion_rois = tumour_rois.squeeze().unsqueeze(0)\n",
    "            background_rois = background.squeeze().unsqueeze(0)\n",
    "        else:\n",
    "            lesion_rois = torch.zeros_like(refs).unsqueeze(0)\n",
    "            background_rois = torch.zeros_like(refs).unsqueeze(0)\n",
    "            lesion_rois[refs.unsqueeze(0)!=0] = 1\n",
    "            background_rois[refs.unsqueeze(0)!=0] = 1\n",
    "        psnrs = []\n",
    "        ssims = []\n",
    "        crcs = []\n",
    "        stds = []\n",
    "        for img_idx in range(images.shape[0]):\n",
    "            image = images[img_idx].cpu().numpy()\n",
    "            ref = refs[img_idx].squeeze().cpu().numpy()\n",
    "            lesion_roi = lesion_rois[img_idx].cpu().numpy()\n",
    "            background_roi = background_rois[img_idx].squeeze().cpu().numpy()\n",
    "            psnr_r = []\n",
    "            ssim_r = []\n",
    "            crc_r = []\n",
    "            b_bar_r = []\n",
    "            for realisation in image:\n",
    "                psnr_r.append(torch.asarray(PSNR(realisation,ref)))\n",
    "                ssim_r.append(torch.asarray(SSIM(realisation, ref)))\n",
    "                if background_roi.sum() != 0:\n",
    "                    background_idx = np.nonzero(background_roi)\n",
    "                    b_bar = realisation[background_idx]\n",
    "                    b_t = ref[background_idx]\n",
    "                    crc_t = []\n",
    "                    for i in range(len(lesion_roi)):\n",
    "                        if lesion_roi[i,:,:].sum() != 0:\n",
    "                            tumour_roi_idx = np.nonzero(lesion_roi[i,:,:])\n",
    "                            a_bar = realisation[tumour_roi_idx]\n",
    "                            a_t = ref[tumour_roi_idx]\n",
    "                            if a_bar.mean() == 0 and b_bar.mean() == 0:\n",
    "                                crc_t.append(np.array([0.0]))\n",
    "                            else:\n",
    "                                crc_t.append((a_bar.mean()/b_bar.mean() - 1) / (a_t.mean()/b_t.mean() - 1))\n",
    "                    crc_r.append(torch.asarray(crc_t).mean())\n",
    "                    b_bar_r.append(torch.asarray(b_bar))\n",
    "            std = (torch.std(torch.stack(b_bar_r), dim=0)/torch.clamp(torch.stack(b_bar_r).mean(0),1e-9)).mean()\n",
    "            psnrs.append(torch.asarray(psnr_r))\n",
    "            ssims.append(torch.asarray(ssim_r))\n",
    "            crcs.append(torch.asarray(crc_r))\n",
    "            stds.append(torch.asarray(std))\n",
    "        stdf.append(torch.stack(stds).mean())\n",
    "        crcf.append(torch.stack(crcs).mean())\n",
    "        psnrf.append(torch.stack(psnrs).mean())\n",
    "        ssimf.append(torch.stack(ssims).mean())\n",
    "        \"\"\" print(torch.stack(psnrs).mean(), torch.stack(ssims).mean(), torch.stack(crcs).mean(), std)\n",
    "        fig, ax = plt.subplots(1,2)\n",
    "        fig.colorbar(ax[0].imshow(x_pred[0,0].cpu().numpy()))\n",
    "        fig.colorbar(ax[1].imshow(ref))\n",
    "        plt.show() \"\"\"\n",
    "    print(torch.stack(psnrf).mean(), torch.stack(ssimf).mean(), torch.stack(crcf).mean(), torch.stack(stdf).mean(), torch.stack(kl_divf).mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "from glob import glob\n",
    "import torch\n",
    "import omegaconf\n",
    "from normalisation import Normalisation\n",
    "from lpd_modules import LPDForwardFunction2D, LPDAdjointFunction2D\n",
    "from lpd import get_lpd_model\n",
    "from unet import get_unet_model\n",
    "from src import PSNR, SSIM\n",
    "\n",
    "detector_efficiency = 1./30\n",
    "\n",
    "models = [\"data_corrected_mean\"]\n",
    "\n",
    "model_base_path = \"/home/user/sirf/results/lpd_unet/LPD/\"\n",
    "save_idxs = [10,20,30,40,50,60]\n",
    "for model_type in models:\n",
    "    model_path = model_base_path + model_type + \"/\"\n",
    "    config = omegaconf.OmegaConf.load(model_path + \".hydra/config.yaml\")\n",
    "    model = get_lpd_model(n_iter = config.benchmark.n_iter, op = LPDForwardFunction2D, op_adj = LPDAdjointFunction2D)\n",
    "    model.load_state_dict(torch.load(sorted(glob(model_path+\"/*/model_min_val_loss.pt\"))[-1]))\n",
    "    model.eval()\n",
    "    model.to(config.device)\n",
    "    print(sum(p.numel() for p in model.parameters()))\n",
    "    get_normalisation = Normalisation(config.benchmark.normalisation)\n",
    "    stdf = []\n",
    "    crcf = []\n",
    "    psnrf = []\n",
    "    ssimf = []\n",
    "    kl_divf = []\n",
    "    for idx, batch in enumerate(test_loader):\n",
    "        ref = torch.swapaxes(batch[0], 0, 1).to(config.device)\n",
    "        osem = torch.swapaxes(batch[2], 0, 1).to(config.device)\n",
    "        measurements = torch.swapaxes(batch[4], 0, 1).to(config.device)\n",
    "        contamination_factor = torch.swapaxes(batch[5], 0, 1)[:,[0],None].to(config.device)[:,0,:]\n",
    "        attn_factors = torch.swapaxes(batch[6], 0, 1).to(config.device)\n",
    "        if len(batch) > 7:\n",
    "            background = torch.swapaxes(batch[7], 0, 1).to(config.device)\n",
    "            tumour_rois = torch.swapaxes(batch[8], 0, 1).to(config.device)\n",
    "        norm = get_normalisation(osem, measurements, contamination_factor)\n",
    "        x_pred = torch.clamp(model(osem, measurements, acq_model, attn_factors/30, norm, contamination_factor).detach(),0)\n",
    "        kldiv_r = kl_div(x = x_pred,\n",
    "\t\t\tacq_model=acq_model, \n",
    "\t\t\tattn_factors=attn_factors, \n",
    "\t\t\tcontamination=contamination_factor.unsqueeze(1), \n",
    "\t\t\tmeasurements=measurements,\n",
    "\t\t\tscale_factor=1.)[0].squeeze()\n",
    "        kl_divf.append(kldiv_r.mean())\n",
    "        \n",
    "        images = x_pred.squeeze().unsqueeze(0)\n",
    "        refs = ref.squeeze().unsqueeze(0)\n",
    "        if len(batch) > 7:\n",
    "            lesion_rois = tumour_rois.squeeze().unsqueeze(0)\n",
    "            background_rois = background.squeeze().unsqueeze(0)\n",
    "        else:\n",
    "            lesion_rois = torch.zeros_like(refs).unsqueeze(0)\n",
    "            background_rois = torch.zeros_like(refs).unsqueeze(0)\n",
    "            lesion_rois[refs.unsqueeze(0)!=0] = 1\n",
    "            background_rois[refs.unsqueeze(0)!=0] = 1\n",
    "        psnrs = []\n",
    "        ssims = []\n",
    "        crcs = []\n",
    "        stds = []\n",
    "        for img_idx in range(images.shape[0]):\n",
    "            image = images[img_idx].cpu().numpy()\n",
    "            ref = refs[img_idx].squeeze().cpu().numpy()\n",
    "            lesion_roi = lesion_rois[img_idx].cpu().numpy()\n",
    "            background_roi = background_rois[img_idx].squeeze().cpu().numpy()\n",
    "            psnr_r = []\n",
    "            ssim_r = []\n",
    "            crc_r = []\n",
    "            b_bar_r = []\n",
    "            for realisation in image:\n",
    "                psnr_r.append(torch.asarray(PSNR(realisation,ref)))\n",
    "                ssim_r.append(torch.asarray(SSIM(realisation, ref)))\n",
    "                if background_roi.sum() != 0:\n",
    "                    background_idx = np.nonzero(background_roi)\n",
    "                    b_bar = realisation[background_idx]\n",
    "                    b_t = ref[background_idx]\n",
    "                    crc_t = []\n",
    "                    for i in range(len(lesion_roi)):\n",
    "                        if lesion_roi[i,:,:].sum() != 0:\n",
    "                            tumour_roi_idx = np.nonzero(lesion_roi[i,:,:])\n",
    "                            a_bar = realisation[tumour_roi_idx]\n",
    "                            a_t = ref[tumour_roi_idx]\n",
    "                            if a_bar.mean() == 0 and b_bar.mean() == 0:\n",
    "                                crc_t.append(np.array([0.0]))\n",
    "                            else:\n",
    "                                crc_t.append((a_bar.mean()/b_bar.mean() - 1) / (a_t.mean()/b_t.mean() - 1))\n",
    "                    crc_r.append(torch.asarray(crc_t).mean())\n",
    "                    b_bar_r.append(torch.asarray(b_bar))\n",
    "            std = (torch.std(torch.stack(b_bar_r), dim=0)/torch.clamp(torch.stack(b_bar_r).mean(0),1e-9)).mean()\n",
    "            psnrs.append(torch.asarray(psnr_r))\n",
    "            ssims.append(torch.asarray(ssim_r))\n",
    "            crcs.append(torch.asarray(crc_r))\n",
    "            stds.append(torch.asarray(std))\n",
    "        stdf.append(torch.stack(stds).mean())\n",
    "        crcf.append(torch.stack(crcs).mean())\n",
    "        psnrf.append(torch.stack(psnrs).mean())\n",
    "        ssimf.append(torch.stack(ssims).mean())\n",
    "        \"\"\" print(torch.stack(psnrs).mean(), torch.stack(ssims).mean(), torch.stack(crcs).mean(), std)\n",
    "        fig, ax = plt.subplots(1,2)\n",
    "        fig.colorbar(ax[0].imshow(x_pred[0,0].cpu().numpy()))\n",
    "        fig.colorbar(ax[1].imshow(ref))\n",
    "        plt.show() \"\"\"\n",
    "    print(torch.stack(psnrf).mean(), torch.stack(ssimf).mean(), torch.stack(crcf).mean(), torch.stack(stdf).mean(), torch.stack(kl_divf).mean())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
